package com.easy.slurm.slurm.Jsch;

import com.alibaba.fastjson.JSONObject;
import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.multipart.MultipartFile;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author bing.bai
 * @create 2025/5/30
 */
@Slf4j
public class JschClient {

    private static final JschClient INSTANCE = new JschClient();
    private static final int DEFAULT_TIMEOUT = 30; // 默认超时时间（秒）
    private static final int THREAD_TIMEOUT = 10; // 线程等待超时时间（秒）

    private final AtomicBoolean inited = new AtomicBoolean(false);
    private Session session;
    private ExecutorService executorService;

    private JschClient() {
    }

    public static JschClient getInstance() {
        return INSTANCE;
    }

    public synchronized void init(String host, String user, String password) throws JSchException {
        if (inited.get()) {
            log.warn("JschClient already initialized");
            return;
        }

        JSch jsch = new JSch();
        session = jsch.getSession(user, host, 22);
        session.setPassword(password);
        session.setConfig("StrictHostKeyChecking", "no");
        session.setServerAliveInterval(30 * 1000);
        session.connect();

        // 创建专用线程池用于流读取
        executorService = Executors.newFixedThreadPool(2);
        inited.set(true);
        log.info("JschClient initialized successfully for host: {}", host);
    }

    public RunCommandResult runCommand(String command) throws JSchException, IOException {
        log.info("run command :{}", command);
        return runCommand(command, DEFAULT_TIMEOUT);
    }

    public void uploadFile(MultipartFile file, String remoteDir, String remoteFileName) {
        // 1. 创建 SFTP 通道
        ChannelSftp channel = null;
        try (InputStream inputStream = file.getInputStream()) {
            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();
            // 2. 检查并创建远程目录（如果不存在）
            try {
                channel.cd(remoteDir);
            } catch (SftpException e) {
                e.printStackTrace();
                // 目录不存在，逐级创建
                String[] folders = remoteDir.split("/");
                StringBuilder path = new StringBuilder();
                for (String folder : folders) {
                    if (folder.isEmpty()) continue;
                    path.append("/").append(folder);
                    try {
                        channel.cd(path.toString());
                    } catch (SftpException ex) {
                        channel.mkdir(path.toString());
                        channel.cd(path.toString());
                        ex.printStackTrace();
                    }
                }
            }

            // 3. 上传文件
            channel.put(inputStream, remoteFileName);
            log.info("文件上传成功: " + remoteDir + "/" + remoteFileName);
        } catch (Exception e) {
            log.error("文件上传失败");
            e.printStackTrace();
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
        }
    }

    public RunCommandResult runCommand(String command, int timeoutSeconds) throws JSchException, IOException {
        if (!inited.get()) {
            throw new IllegalStateException("JschClient not initialized");
        }

        RunCommandResult result = new RunCommandResult();
        ChannelExec channel = null;

        try {
            channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(command);
            channel.setInputStream(null);

            // 准备读取流
            InputStream in = channel.getInputStream();
            InputStream err = channel.getErrStream();

            channel.connect();

            // 使用CompletableFuture并行读取两个流
            CompletableFuture<List<String>> stdoutFuture = readStreamAsync(in, "stdout");
            CompletableFuture<List<String>> stderrFuture = readStreamAsync(err, "stderr");

            // 等待流读取完成或超时
            CompletableFuture<Void> allFutures = CompletableFuture.allOf(stdoutFuture, stderrFuture);
            try {
                allFutures.get(timeoutSeconds, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                log.warn("Command execution interrupted", e);
            } catch (ExecutionException e) {
                log.error("Error reading command output", e.getCause());
            } catch (TimeoutException e) {
                log.warn("Command execution timed out after {} seconds", timeoutSeconds);
                // 取消未完成的读取任务
                stdoutFuture.cancel(true);
                stderrFuture.cancel(true);
            }

            // 获取结果（即使超时也获取已读取的部分）
            result.setStdout(stdoutFuture.getNow(new ArrayList<>()));
            result.setStderr(stderrFuture.getNow(new ArrayList<>()));
            result.setExitStatus(channel.getExitStatus());
            System.out.println(JSONObject.toJSONString(result));
            return result;
        } finally {
            if (channel != null) {
                channel.disconnect();
            }
        }
    }

    private CompletableFuture<List<String>> readStreamAsync(InputStream stream, String streamType) {
        return CompletableFuture.supplyAsync(() -> {
            List<String> lines = new ArrayList<>();
            try (BufferedReader reader = new BufferedReader(
                    new InputStreamReader(stream, StandardCharsets.UTF_8))) {
                String line;
                while ((line = reader.readLine()) != null) {
                    lines.add(line);
                    // 检查线程是否被中断（超时或取消时）
                    if (Thread.currentThread().isInterrupted()) {
                        log.debug("{} reading interrupted", streamType);
                        break;
                    }
                }
            } catch (IOException e) {
                log.error("Error reading {} stream", streamType, e);
            }
            return lines;
        }, executorService);
    }

    public synchronized void destroy() {
        if (!inited.get()) {
            log.warn("JschClient not initialized, nothing to destroy");
            return;
        }

        if (session != null && session.isConnected()) {
            session.disconnect();
        }

        if (executorService != null) {
            try {
                executorService.shutdown();
                if (!executorService.awaitTermination(THREAD_TIMEOUT, TimeUnit.SECONDS)) {
                    executorService.shutdownNow();
                }
            } catch (InterruptedException e) {
                executorService.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }

        inited.set(false);
        log.info("JschClient destroyed successfully");
    }

    public static class RunCommandResult {
        private List<String> stdout = new ArrayList<>();
        private List<String> stderr = new ArrayList<>();
        private int exitStatus;

        // Getters and setters
        public List<String> getStdout() {
            return stdout;
        }

        public void setStdout(List<String> stdout) {
            this.stdout = stdout;
        }

        public List<String> getStderr() {
            return stderr;
        }

        public void setStderr(List<String> stderr) {
            this.stderr = stderr;
        }

        public int getExitStatus() {
            return exitStatus;
        }

        public void setExitStatus(int exitStatus) {
            this.exitStatus = exitStatus;
        }

        public boolean isSuccess() {
            return exitStatus == 0;
        }

        public String getErrorMsg() {
            String errorMsg = "";
            for (int i = 0; i < stderr.size(); i++) {
                errorMsg += stderr.get(i);
            }
            return errorMsg;
        }

        public String getSuccessMsg() {
            String successMsg = "";
            for (int i = 0; i < stdout.size(); i++) {
                successMsg += stdout.get(i);
            }
            return successMsg;
        }

    }
}